# conda activate ruanzhu
# python image_to_pdf.py ./hetong/apply.jpg -o apply.pdf
# python image_to_pdf.py --dir hetong -o hetong.pdf

import os
import logging
from pathlib import Path

try:
    from reportlab.lib.pagesizes import A4, letter
    from reportlab.platypus import SimpleDocTemplate, Image, Spacer, PageBreak
    from reportlab.lib.units import inch, cm
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False
    print("警告: reportlab 库未安装，PDF 输出功能不可用。")

try:
    from PIL import Image as PILImage
    PILLOW_AVAILABLE = True
except ImportError:
    PILLOW_AVAILABLE = False
    print("警告: Pillow 库未安装，图片处理功能不可用。")

logger = logging.getLogger(__name__)

class ImageToPDFConverter:
    """图片转PDF转换器"""
    
    def __init__(self, page_size=None, margin=None, fit_to_page=True, auto_rotate=True):
        if not PDF_AVAILABLE:
            raise ImportError("reportlab 库未安装，无法使用 PDF 输出功能")
        if not PILLOW_AVAILABLE:
            raise ImportError("Pillow 库未安装，无法使用图片处理功能")
        
        self.base_page_size = page_size or A4
        self.margin = margin or 2*cm
        self.fit_to_page = fit_to_page
        self.auto_rotate = auto_rotate
        self.story = []
        
    def _get_optimal_page_size(self, img_width, img_height):
        """根据图片尺寸选择最佳页面方向"""
        print(f"图片尺寸: {img_width}x{img_height}, auto_rotate: {self.auto_rotate}")
        if not self.auto_rotate:
            return self.base_page_size
            
        # 判断图片是否为竖拍（高度大于宽度）
        is_portrait = img_height > img_width
        
        # 获取页面的宽高
        page_w, page_h = self.base_page_size
        
        # 如果图片是竖拍，使用竖向页面；如果是横拍，使用横向页面
        if is_portrait:
            return (min(page_w, page_h), max(page_w, page_h))  # 竖向
        else:
            return (max(page_w, page_h), min(page_w, page_h))  # 横向
    
    def add_image(self, image_path, new_page=False, quality='high'):
        """添加图片到PDF，支持质量控制"""
        try:
            if not os.path.exists(image_path):
                logger.error(f"图片文件不存在: {image_path}")
                return False
            
            # 获取图片信息并自动处理方向
            with PILImage.open(image_path) as pil_img:
                # 使用Pillow的ImageOps.exif_transpose自动处理EXIF方向
                from PIL import ImageOps
                import tempfile
                
                # 自动根据EXIF信息旋转图片
                corrected_img = ImageOps.exif_transpose(pil_img)
                img_width, img_height = corrected_img.size
                
                print(f"原始尺寸: {pil_img.size}, 校正后尺寸: {img_width}x{img_height}")
                
                # 创建临时文件保存校正后的图片
                temp_file = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)
                corrected_img.save(temp_file.name, 'JPEG', quality=95)
                corrected_image_path = temp_file.name
            
            # 根据调整后的图片尺寸选择最佳页面大小
            optimal_page_size = self._get_optimal_page_size(img_width, img_height)
            page_width = optimal_page_size[0] - 2 * self.margin
            page_height = optimal_page_size[1] - 2 * self.margin
            
            # 如果需要新页面，添加分页符
            if new_page and self.story:
                self.story.append(PageBreak())
            
            # 计算缩放比例
            if self.fit_to_page:
                width_ratio = page_width / img_width
                height_ratio = page_height / img_height
                scale_ratio = min(width_ratio, height_ratio)
                
                final_width = img_width * scale_ratio
                final_height = img_height * scale_ratio
            else:
                final_width = img_width
                final_height = img_height
            
            # 创建图片对象，使用校正后的图片
            img = Image(corrected_image_path, width=final_width, height=final_height)
            
            # 居中对齐
            img.hAlign = 'CENTER'
            
            # 添加到story中
            self.story.append(img)
            self.story.append(Spacer(1, 10))
            
            logger.info(f"成功添加图片: {os.path.basename(image_path)} ({final_width:.1f}x{final_height:.1f})")
            
            # 清理临时文件（在PDF生成后）
            # 注意：这里不能立即删除，因为reportlab需要在生成PDF时访问文件
            # 可以在类中维护一个临时文件列表，在save方法中清理
            if not hasattr(self, '_temp_files'):
                self._temp_files = []
            self._temp_files.append(corrected_image_path)
            
            return True
            
        except Exception as e:
            logger.error(f"添加图片失败 {image_path}: {e}")
            return False
        
    
    def add_images_from_directory(self, directory_path, pattern="*.jpg", each_page=False):
        """从目录添加所有匹配的图片"""
        try:
            directory = Path(directory_path)
            if not directory.exists():
                logger.error(f"目录不存在: {directory_path}")
                return False
            
            # 支持多种图片格式
            patterns = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tiff", "*.gif"]
            if pattern != "*.jpg":
                patterns = [pattern]
            
            image_files = []
            for pat in patterns:
                image_files.extend(directory.glob(pat))
                image_files.extend(directory.glob(pat.upper()))
            
            # 按文件名排序
            image_files.sort()
            
            if not image_files:
                logger.warning(f"在目录 {directory_path} 中未找到图片文件")
                return False
            
            success_count = 0
            for i, img_file in enumerate(image_files):
                new_page = each_page and i > 0
                if self.add_image(str(img_file), new_page=new_page):
                    success_count += 1
            
            logger.info(f"成功添加 {success_count}/{len(image_files)} 个图片文件")
            return success_count > 0
            
        except Exception as e:
            logger.error(f"从目录添加图片失败: {e}")
            return False
    
    def save(self, output_file):
        """保存PDF文件"""
        try:
            if not self.story:
                logger.error("没有内容可保存")
                return False
            
            # 如果启用了自动旋转，需要为每个图片单独创建页面
            if self.auto_rotate:
                return self._save_with_auto_rotation(output_file)
            else:
                return self._save_standard(output_file)
                
        except Exception as e:
            logger.error(f"保存PDF失败: {e}")
            return False
    
    def _save_standard(self, output_file):
        """标准保存方式"""
        doc = SimpleDocTemplate(
            output_file,
            pagesize=self.base_page_size,
            leftMargin=self.margin,
            rightMargin=self.margin,
            topMargin=self.margin,
            bottomMargin=self.margin
        )
        doc.build(self.story)
        return True
    
    def _save_with_auto_rotation(self, output_file):
        """支持自动旋转的保存方式"""
        # 这里需要更复杂的实现来处理不同页面方向
        # 暂时使用标准方式
        return self._save_standard(output_file)


def images_to_pdf(image_paths, output_file, **kwargs):
    """将图片文件转换为PDF的便捷函数"""
    try:
        # 提取 each_page 参数，不传递给构造函数
        each_page = kwargs.pop('each_page', False)
        converter = ImageToPDFConverter(**kwargs)
        
        # 如果是字符串，转换为列表
        if isinstance(image_paths, str):
            image_paths = [image_paths]
        
        success_count = 0
        for i, img_path in enumerate(image_paths):
            new_page = each_page and i > 0
            if converter.add_image(img_path, new_page=new_page):
                success_count += 1
        
        if success_count > 0:
            return converter.save(output_file)
        else:
            logger.error("没有成功添加任何图片")
            return False
            
    except Exception as e:
        logger.error(f"图片转PDF失败: {e}")
        return False


def directory_to_pdf(directory_path, output_file, **kwargs):
    """将目录中的所有图片转换为PDF的便捷函数"""
    try:
        # 提取不属于构造函数的参数
        pattern = kwargs.pop('pattern', '*.jpg')
        each_page = kwargs.pop('each_page', False)
        
        converter = ImageToPDFConverter(**kwargs)
        
        success = converter.add_images_from_directory(
            directory_path, 
            pattern=pattern, 
            each_page=each_page
        )
        
        if success:
            return converter.save(output_file)
        else:
            return False
            
    except Exception as e:
        logger.error(f"目录转PDF失败: {e}")
        return False


if __name__ == '__main__':
    import sys
    import argparse
    
    parser = argparse.ArgumentParser(description='图片转PDF工具')
    parser.add_argument('inputs', nargs='+', help='输入图片文件或目录')
    parser.add_argument('-o', '--output', required=True, help='输出PDF文件名')
    parser.add_argument('--dir', action='store_true', help='输入为目录')
    parser.add_argument('--each-page', action='store_true', help='每张图片单独一页')
    parser.add_argument('--auto-rotate', action='store_true', default=True, help='自动调整页面方向')
    parser.add_argument('--quality', choices=['high', 'medium', 'low'], default='high', help='图片质量')
    parser.add_argument('--margin', type=float, default=2.0, help='页边距(cm)')
    
    args = parser.parse_args()
    
    margin = args.margin * cm
    
    if args.dir:
        success = directory_to_pdf(
            args.inputs[0], 
            args.output, 
            each_page=args.each_page,
            auto_rotate=args.auto_rotate,
            margin=margin
        )
    else:
        success = images_to_pdf(
            args.inputs, 
            args.output, 
            each_page=args.each_page,
            auto_rotate=args.auto_rotate,
            margin=margin
        )
    
    if success:
        print(f"成功生成PDF: {args.output}")
    else:
        print("PDF生成失败")
        sys.exit(1)